#!/usr/bin/env julia
using DelimitedFiles
using Printf
using Plots

agg_path = "results/drp4dvar_vs_3dvar/metrics_aggregated.csv"
rows = readdlm(agg_path, ',', header=true)[1]

ensembles = Int[]
cost_drp = Float64[]
cost_3dvar = Float64[]
rms_drp = Float64[]
solve_drp = Float64[]
rms_bg = Float64[]

for row in eachrow(rows)
    method = String(row[2])
    if method == "background"
        push!(rms_bg, row[4])
    elseif method == "drp4dvar"
        push!(ensembles, Int(round(row[7])))
        push!(cost_drp, row[3])
        push!(rms_drp, row[4])
        push!(solve_drp, row[6])
    elseif method == "3dvar"
        push!(cost_3dvar, row[3])
    end
end

sorted = sortperm(ensembles)
ens_sorted = ensembles[sorted]
cost_sorted = cost_drp[sorted]
rms_sorted = rms_drp[sorted]
solve_sorted = solve_drp[sorted]

bg_rms = isempty(rms_bg) ? NaN : rms_bg[1]
three_dvar_cost = isempty(cost_3dvar) ? NaN : cost_3dvar[1]

plot_dir = "results/drp4dvar_vs_3dvar"
mkpath(plot_dir)

# Cost plot
p1 = plot(ens_sorted, cost_sorted, marker=:o, xlabel="Ensemble Size", ylabel="Cost", title="DRP Cost vs Ensemble", label="DRP cost")
if !isnan(three_dvar_cost)
    hline!(p1, [three_dvar_cost], linestyle=:dash, label="3DVar cost")
end
savefig(p1, joinpath(plot_dir, "drp_cost_vs_ensemble.png"))

# RMS innovation plot
p2 = plot(ens_sorted, rms_sorted, marker=:o, xlabel="Ensemble Size", ylabel="RMS Innovation", title="DRP Innovation RMS", label="DRP")
if !isnan(bg_rms)
    hline!(p2, [bg_rms], linestyle=:dash, label="Background")
end
savefig(p2, joinpath(plot_dir, "drp_rms_vs_ensemble.png"))

# Solve time plot
p3 = plot(ens_sorted, solve_sorted, marker=:o, xlabel="Ensemble Size", ylabel="Solve Time (s)", title="DRP Solve Time", label="DRP")
savefig(p3, joinpath(plot_dir, "drp_time_vs_ensemble.png"))

println("Generated DRP comparison plots in $plot_dir")
